
#include <windows.h>
#include <wtsapi32.h>
#include <tchar.h>
#include <stdio.h>
#include <iostream>
#include <vcclr.h>


#include "Stdafx.h"
#include "SessionLauncher.h"

using namespace System;

namespace Sustainalytics{
	namespace OS {

		LinkedProcess^ LinkedProcess::Start(System::String^ cmdWithParamenters, bool spawnInConsoleSession, bool hidden)
		{
			wchar_t * appName = nullptr;
			wchar_t *cmd = (wchar_t*)ToUnmanged(cmdWithParamenters);

			HANDLE job = CreateJobObject(nullptr, nullptr);
			if (!job) {
				throw  CreateManagedException(1, GetLastError());

			}

			HANDLE ioPort = CreateIoCompletionPort(INVALID_HANDLE_VALUE, nullptr, 0, 1);
			if (!ioPort) {
				throw  CreateManagedException(1, GetLastError());
			}

			JOBOBJECT_ASSOCIATE_COMPLETION_PORT port;
			port.CompletionKey = job;
			port.CompletionPort = ioPort;
			if (!SetInformationJobObject(job, JobObjectAssociateCompletionPortInformation, &port, sizeof(port))) {
				throw  CreateManagedException(1, GetLastError());
			}


			JOBOBJECT_BASIC_LIMIT_INFORMATION basicLimit;
			basicLimit.LimitFlags = JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE;
			JOBOBJECT_EXTENDED_LIMIT_INFORMATION extendedLimit;
			extendedLimit.BasicLimitInformation = basicLimit;

			if (!SetInformationJobObject(job, JobObjectExtendedLimitInformation, &extendedLimit, sizeof(extendedLimit))) {
				throw  CreateManagedException(1, GetLastError());
			}


			PROCESS_INFORMATION processInformation;
			STARTUPINFO startupInfo = { sizeof(startupInfo) };
			DWORD dwCreationFlags = CREATE_SUSPENDED;
			if (hidden){
				startupInfo.dwFlags = STARTF_USESHOWWINDOW;
				startupInfo.wShowWindow = SW_HIDE;
			}
			else {
				wchar_t *dsktp = L"WinSta0\\Default"; //allow process to be interactive
				startupInfo.lpDesktop = dsktp;

				dwCreationFlags = dwCreationFlags | CREATE_NEW_CONSOLE;
			}

			if (!spawnInConsoleSession) {
				if (!CreateProcess(appName, cmd, nullptr, nullptr, FALSE, dwCreationFlags, nullptr, nullptr, &startupInfo, &processInformation)) {
					throw  CreateManagedException(1, GetLastError());
				}
			}
			else {
				HANDLE svctoken = nullptr;
				HANDLE provtoken = nullptr;
				DWORD sessionId = WTSGetActiveConsoleSessionId();

				// Create token for new process by duplicating our own token
				if (OpenProcessToken(GetCurrentProcess(), TOKEN_ALL_ACCESS, &svctoken) == 0)
				{
					throw  CreateManagedException(1, GetLastError());
				}
				DuplicateTokenEx(svctoken, TOKEN_ALL_ACCESS, nullptr, SecurityImpersonation, TokenPrimary, &provtoken);
				CloseHandle(svctoken);
				svctoken = nullptr;
				// Move the new token to the target session
				if (!SetTokenInformation(provtoken, TokenSessionId, &sessionId, sizeof(DWORD)))
				{
					throw  CreateManagedException(1, GetLastError());
				}

				// Set up the new process information
				if (!CreateProcessAsUser(provtoken, appName, cmd, nullptr, nullptr, FALSE, CREATE_SUSPENDED, nullptr, nullptr, &startupInfo, &processInformation))
				{
					throw  CreateManagedException(1, GetLastError());
				}

				if (provtoken)
				{
					CloseHandle(provtoken);
				}
			}
			//if this fails the spawned process remain suspended !!!!
			if (!AssignProcessToJobObject(job, processInformation.hProcess)) {
				throw  CreateManagedException(1, GetLastError());
			}

			ResumeThread(processInformation.hThread);

			CloseHandle(processInformation.hThread);
			CloseHandle(processInformation.hProcess);

			return gcnew LinkedProcess(ioPort, job, processInformation.dwProcessId);
		}


		ProcessLaunchException^ LinkedProcess::CreateManagedException(int errorCode, int win32ErrorCode){

			ProcessLaunchException^ exception = gcnew ProcessLaunchException(errorCode, gcnew System::ComponentModel::Win32Exception(win32ErrorCode));
			return exception;
		}
	}
}